package job_manager

import (
	"context"
	"errors"
	"fmt"
	"log"
	"sync"
	"time"

	"gitlab.ctyuncdn.cn/tai/infra/cloud-controller/pkg/mysql"
)

// JobState represents the state of a job
type JobState string

const (
	JobStatePending   JobState = "pending"
	JobStateRunning   JobState = "running"
	JobStateCompleted JobState = "completed"
	JobStateFailed    JobState = "failed"
	JobStateCancelled JobState = "cancelled"
)

// JobType represents the type of a job
type JobType string

const (
	JobTypeNotebook  JobType = "notebook"
	JobTypeTraining  JobType = "training"
	JobTypeInference JobType = "inference"
	JobTypeBatch     JobType = "batch"
)

// Error definitions
var (
	ErrUnsupportedJobType = errors.New("unsupported job type")
)

// JobManager defines the interface for job management
type JobManager interface {
	// Core job management
	CreateJob(ctx context.Context, job *mysql.Job) error
	StartJob(ctx context.Context, jobID string) error
	StopJob(ctx context.Context, jobID string) error
	CancelJob(ctx context.Context, jobID string) error
	GetJob(ctx context.Context, jobID string) (*mysql.Job, error)
	ListJobs(ctx context.Context, filters map[string]interface{}) ([]*mysql.Job, error)

	// Job lifecycle
	UpdateJobStatus(ctx context.Context, jobID string, status JobState, message string) error
	GetJobStatus(ctx context.Context, jobID string) (*mysql.JobStatus, error)
	GetJobHistory(ctx context.Context, jobID string) ([]*mysql.JobStatus, error)

	// Resource management
	GetJobResources(ctx context.Context, jobID string) ([]*Resource, error)
	CreateJobResources(ctx context.Context, jobID string) error
	DeleteJobResources(ctx context.Context, jobID string) error

	// Job type management
	RegisterJobType(jobType JobType, handler JobTypeHandler) error
	GetJobTypeHandler(jobType JobType) (JobTypeHandler, error)

	// Health and monitoring
	IsHealthy() bool
	GetStats() map[string]interface{}
}

// JobTypeHandler defines the interface for job type-specific operations
type JobTypeHandler interface {
	// Job type validation
	ValidateJob(job *mysql.Job) error
	ValidateParameters(params map[string]interface{}) error

	// Resource management
	CreateResources(ctx context.Context, job *mysql.Job) ([]*Resource, error)
	DeleteResources(ctx context.Context, job *mysql.Job) error
	GetResourceTemplates(job *mysql.Job) ([]*ResourceTemplate, error)

	// Job execution
	StartJob(ctx context.Context, job *mysql.Job) error
	StopJob(ctx context.Context, job *mysql.Job) error
	GetJobStatus(ctx context.Context, job *mysql.Job) (*mysql.JobStatus, error)

	// Job type info
	GetJobType() JobType
	GetDescription() string
	GetSupportedParameters() []Parameter
}

// Resource represents a Kubernetes resource
type Resource struct {
	Type      string                 `json:"type"`
	Name      string                 `json:"name"`
	Namespace string                 `json:"namespace"`
	Spec      map[string]interface{} `json:"spec"`
	Status    map[string]interface{} `json:"status"`
	CreatedAt time.Time              `json:"created_at"`
	UpdatedAt time.Time              `json:"updated_at"`
}

// ResourceTemplate represents a template for creating resources
type ResourceTemplate struct {
	Type        string                 `json:"type"`
	Name        string                 `json:"name"`
	Namespace   string                 `json:"namespace"`
	Template    map[string]interface{} `json:"template"`
	Parameters  []Parameter            `json:"parameters"`
	Description string                 `json:"description"`
}

// Parameter represents a job parameter
type Parameter struct {
	Name        string      `json:"name"`
	Type        string      `json:"type"`
	Required    bool        `json:"required"`
	Default     interface{} `json:"default"`
	Description string      `json:"description"`
	Validation  string      `json:"validation"`
}

// BaseJobManager provides a base implementation of JobManager
type BaseJobManager struct {
	db       DatabaseInterface
	jobTypes map[JobType]JobTypeHandler
	mu       sync.RWMutex
	ctx      context.Context
	cancel   context.CancelFunc
	wg       sync.WaitGroup
}

// DatabaseInterface defines the interface for database operations
type DatabaseInterface interface {
	Create(value interface{}) DatabaseResult
	Model(value interface{}) DatabaseInterface
	Where(query interface{}, args ...interface{}) DatabaseInterface
	Order(value interface{}) DatabaseInterface
	First(dest interface{}) DatabaseResult
	Find(dest interface{}) DatabaseResult
	Updates(values interface{}) DatabaseResult
	Error() error
}

// DatabaseResult represents the result of a database operation
type DatabaseResult interface {
	Error() error
}

// NewBaseJobManager creates a new base job manager
func NewBaseJobManager(db DatabaseInterface) *BaseJobManager {
	ctx, cancel := context.WithCancel(context.Background())

	return &BaseJobManager{
		db:       db,
		jobTypes: make(map[JobType]JobTypeHandler),
		ctx:      ctx,
		cancel:   cancel,
	}
}

// Start starts the job manager
func (m *BaseJobManager) Start() error {
	log.Println("Job manager started successfully")
	return nil
}

// Stop stops the job manager
func (m *BaseJobManager) Stop() error {
	log.Println("Stopping job manager...")

	m.cancel()
	m.wg.Wait()

	log.Println("Job manager stopped successfully")
	return nil
}

// CreateJob creates a new job
func (m *BaseJobManager) CreateJob(ctx context.Context, job *mysql.Job) error {
	// Validate job
	if err := m.validateJob(job); err != nil {
		return fmt.Errorf("job validation failed: %w", err)
	}

	// Set initial state
	job.Status = string(JobStatePending)
	job.CreatedAt = time.Now()
	job.UpdatedAt = time.Now()

	// Save job to database
	if err := m.db.Create(job).Error(); err != nil {
		return fmt.Errorf("failed to create job: %w", err)
	}

	// Create initial status
	status := &mysql.JobStatus{
		JobID:     job.JobID,
		Status:    string(JobStatePending),
		Message:   "Job created successfully",
		CreatedAt: time.Now(),
	}

	if err := m.db.Create(status).Error(); err != nil {
		return fmt.Errorf("failed to create job status: %w", err)
	}

	log.Printf("Created job: %s", job.JobID)
	return nil
}

// StartJob starts a job
func (m *BaseJobManager) StartJob(ctx context.Context, jobID string) error {
	// Get job
	job, err := m.GetJob(ctx, jobID)
	if err != nil {
		return fmt.Errorf("failed to get job: %w", err)
	}

	// Get job type handler
	handler, err := m.GetJobTypeHandler(JobType(job.JobType))
	if err != nil {
		return fmt.Errorf("failed to get job type handler: %w", err)
	}

	// Create resources
	if err := m.CreateJobResources(ctx, jobID); err != nil {
		return fmt.Errorf("failed to create job resources: %w", err)
	}

	// Start job using handler
	if err := handler.StartJob(ctx, job); err != nil {
		return fmt.Errorf("failed to start job: %w", err)
	}

	// Update job status
	if err := m.UpdateJobStatus(ctx, jobID, JobStateRunning, "Job started successfully"); err != nil {
		return fmt.Errorf("failed to update job status: %w", err)
	}

	log.Printf("Started job: %s", jobID)
	return nil
}

// StopJob stops a job
func (m *BaseJobManager) StopJob(ctx context.Context, jobID string) error {
	// Get job
	job, err := m.GetJob(ctx, jobID)
	if err != nil {
		return fmt.Errorf("failed to get job: %w", err)
	}

	// Get job type handler
	handler, err := m.GetJobTypeHandler(JobType(job.JobType))
	if err != nil {
		return fmt.Errorf("failed to get job type handler: %w", err)
	}

	// Stop job using handler
	if err := handler.StopJob(ctx, job); err != nil {
		return fmt.Errorf("failed to stop job: %w", err)
	}

	// Update job status
	if err := m.UpdateJobStatus(ctx, jobID, JobStateCompleted, "Job stopped successfully"); err != nil {
		return fmt.Errorf("failed to update job status: %w", err)
	}

	log.Printf("Stopped job: %s", jobID)
	return nil
}

// CancelJob cancels a job
func (m *BaseJobManager) CancelJob(ctx context.Context, jobID string) error {
	// Get job
	job, err := m.GetJob(ctx, jobID)
	if err != nil {
		return fmt.Errorf("failed to get job: %w", err)
	}

	// Get job type handler
	handler, err := m.GetJobTypeHandler(JobType(job.JobType))
	if err != nil {
		return fmt.Errorf("failed to get job type handler: %w", err)
	}

	// Stop job using handler
	if err := handler.StopJob(ctx, job); err != nil {
		return fmt.Errorf("failed to stop job: %w", err)
	}

	// Delete resources
	if err := m.DeleteJobResources(ctx, jobID); err != nil {
		return fmt.Errorf("failed to delete job resources: %w", err)
	}

	// Update job status
	if err := m.UpdateJobStatus(ctx, jobID, JobStateCancelled, "Job cancelled successfully"); err != nil {
		return fmt.Errorf("failed to update job status: %w", err)
	}

	log.Printf("Cancelled job: %s", jobID)
	return nil
}

// GetJob gets a job by ID
func (m *BaseJobManager) GetJob(ctx context.Context, jobID string) (*mysql.Job, error) {
	var job mysql.Job
	if err := m.db.Where("job_id = ?", jobID).First(&job).Error(); err != nil {
		return nil, fmt.Errorf("failed to get job: %w", err)
	}

	return &job, nil
}

// ListJobs lists jobs with optional filters
func (m *BaseJobManager) ListJobs(ctx context.Context, filters map[string]interface{}) ([]*mysql.Job, error) {
	query := m.db.Model(&mysql.Job{})

	// Apply filters
	for key, value := range filters {
		query = query.Where(key+" = ?", value)
	}

	var jobs []*mysql.Job
	if err := query.Order("created_at DESC").Find(&jobs).Error(); err != nil {
		return nil, fmt.Errorf("failed to list jobs: %w", err)
	}

	return jobs, nil
}

// UpdateJobStatus updates the status of a job
func (m *BaseJobManager) UpdateJobStatus(ctx context.Context, jobID string, status JobState, message string) error {
	// Update job status
	if err := m.db.Model(&mysql.Job{}).Where("job_id = ?", jobID).Updates(map[string]interface{}{
		"status":     string(status),
		"updated_at": time.Now(),
	}).Error(); err != nil {
		return fmt.Errorf("failed to update job status: %w", err)
	}

	// Create status record
	statusRecord := &mysql.JobStatus{
		JobID:     jobID,
		Status:    string(status),
		Message:   message,
		CreatedAt: time.Now(),
	}

	if err := m.db.Create(statusRecord).Error(); err != nil {
		return fmt.Errorf("failed to create status record: %w", err)
	}

	return nil
}

// GetJobStatus gets the current status of a job
func (m *BaseJobManager) GetJobStatus(ctx context.Context, jobID string) (*mysql.JobStatus, error) {
	var status mysql.JobStatus
	if err := m.db.Where("job_id = ?", jobID).Order("created_at DESC").First(&status).Error(); err != nil {
		return nil, fmt.Errorf("failed to get job status: %w", err)
	}

	return &status, nil
}

// GetJobHistory gets the status history of a job
func (m *BaseJobManager) GetJobHistory(ctx context.Context, jobID string) ([]*mysql.JobStatus, error) {
	var statuses []*mysql.JobStatus
	if err := m.db.Where("job_id = ?", jobID).Order("created_at DESC").Find(&statuses).Error(); err != nil {
		return nil, fmt.Errorf("failed to get job history: %w", err)
	}

	return statuses, nil
}

// GetJobResources gets the resources for a job
func (m *BaseJobManager) GetJobResources(ctx context.Context, jobID string) ([]*Resource, error) {
	// This would typically query Kubernetes or a resource database
	// For now, return empty list
	return []*Resource{}, nil
}

// CreateJobResources creates resources for a job
func (m *BaseJobManager) CreateJobResources(ctx context.Context, jobID string) error {
	// Get job
	job, err := m.GetJob(ctx, jobID)
	if err != nil {
		return fmt.Errorf("failed to get job: %w", err)
	}

	// Get job type handler
	handler, err := m.GetJobTypeHandler(JobType(job.JobType))
	if err != nil {
		return fmt.Errorf("failed to get job type handler: %w", err)
	}

	// Create resources using handler
	_, err = handler.CreateResources(ctx, job)
	if err != nil {
		return fmt.Errorf("failed to create job resources: %w", err)
	}

	return nil
}

// DeleteJobResources deletes resources for a job
func (m *BaseJobManager) DeleteJobResources(ctx context.Context, jobID string) error {
	// Get job
	job, err := m.GetJob(ctx, jobID)
	if err != nil {
		return fmt.Errorf("failed to get job: %w", err)
	}

	// Get job type handler
	handler, err := m.GetJobTypeHandler(JobType(job.JobType))
	if err != nil {
		return fmt.Errorf("failed to get job type handler: %w", err)
	}

	// Delete resources using handler
	if err := handler.DeleteResources(ctx, job); err != nil {
		return fmt.Errorf("failed to delete job resources: %w", err)
	}

	return nil
}

// RegisterJobType registers a job type handler
func (m *BaseJobManager) RegisterJobType(jobType JobType, handler JobTypeHandler) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	if handler == nil {
		return fmt.Errorf("handler cannot be nil")
	}

	m.jobTypes[jobType] = handler
	log.Printf("Registered job type handler: %s", jobType)

	return nil
}

// GetJobTypeHandler gets a job type handler
func (m *BaseJobManager) GetJobTypeHandler(jobType JobType) (JobTypeHandler, error) {
	m.mu.RLock()
	defer m.mu.RUnlock()

	handler, exists := m.jobTypes[jobType]
	if !exists {
		return nil, fmt.Errorf("job type handler not found: %s", jobType)
	}

	return handler, nil
}

// IsHealthy checks if the job manager is healthy
func (m *BaseJobManager) IsHealthy() bool {
	// Simple health check - could be enhanced with more checks
	return true
}

// GetStats returns job manager statistics
func (m *BaseJobManager) GetStats() map[string]interface{} {
	m.mu.RLock()
	defer m.mu.RUnlock()

	stats := map[string]interface{}{
		"registered_job_types": len(m.jobTypes),
		"healthy":              m.IsHealthy(),
	}

	// Add job type information
	jobTypes := make([]string, 0, len(m.jobTypes))
	for jobType := range m.jobTypes {
		jobTypes = append(jobTypes, string(jobType))
	}
	stats["job_types"] = jobTypes

	return stats
}

// validateJob validates a job
func (m *BaseJobManager) validateJob(job *mysql.Job) error {
	if job == nil {
		return fmt.Errorf("job cannot be nil")
	}

	if job.JobID == "" {
		return fmt.Errorf("job_id is required")
	}

	if job.JobType == "" {
		return fmt.Errorf("job_type is required")
	}

	// Check if job type is supported
	_, err := m.GetJobTypeHandler(JobType(job.JobType))
	if err != nil {
		return fmt.Errorf("unsupported job type: %s", job.JobType)
	}

	return nil
}
